import json
import os

import numpy as np
import matplotlib.pyplot as plt

line_styles = [
    '-',  # 实线
    '--',  # 虚线
    '-.',  # 点画线
    ':',  # 点实线
]


def _read_acc(file):
    acc = []
    with open(file, "r", encoding="utf-8") as f:
        lines = f.readlines()
        for line in lines:
            arr = line.strip().split('\t')
            acc.append(float(arr[1]))
    return acc


def moving_average_filter(data, window_size):
    # 定义滑动窗口
    window = np.ones(int(window_size)) / float(window_size)
    # 对数据进行滑动平均滤波
    filtered_data = np.convolve(data, window, mode='valid')
    return filtered_data


def draw_acc(y_datas: dict, save_file="acc.png", x_scale=100, title=""):
    fig, ax = plt.subplots()
    ax.set_xlabel('Epochs', fontsize=20)
    ax.set_ylabel('Test accuracy(%)', fontsize=20)
    ax.set_title(title, fontsize=20)

    ax.set_yticks(np.arange(0, 100, 10))
    ax.set_yticklabels(np.arange(0, 100, 10))
    ax.set_ylim([0, 100])
    ax.set_xlim([0, x_scale])
    ax.grid()

    x_data = np.arange(0, x_scale, 1)
    count = 0
    for name, data in y_datas.items():
        data = [float(i) * 100 for i in data]
        ax.plot(x_data, data[: x_scale], linewidth=2, label=name, linestyle=line_styles[count])
        count += 1

    plt.legend(loc='lower right', fontsize=16)

    save_format = save_file.split(".")[-1]
    fig.savefig(save_file, format=save_format, bbox_inches="tight")
    # plt.show()


def draw_ppl(y_datas: dict, save_file="acc.svg", x_scale=100, title=""):
    fig, ax = plt.subplots()
    ax.set_xlabel('Epochs', fontsize=20)
    ax.set_ylabel('Test perplexity', fontsize=20)
    ax.set_title(title, fontsize=20)

    ax.set_yticks(np.arange(1, 1.3, 0.1))
    # ax.set_yticklabels(np.arange(1, 2, 0.1))
    ax.set_yticklabels([f"{i:.1f}" for i in np.arange(1, 1.3, 0.1)])  # 设置刻度标签
    ax.set_ylim([1, 1.3])
    ax.set_xlim([0, x_scale])
    ax.grid()

    x_data = np.arange(0, x_scale, 1)
    count = 0
    for name, data in y_datas.items():
        data = [float(i) for i in data]
        ax.plot(x_data, data[: x_scale], linewidth=2, label=name, linestyle=line_styles[count])
        count += 1

    plt.legend(loc='upper right', fontsize=16)

    save_format = save_file.split(".")[-1]
    fig.savefig(save_file, format=save_format, bbox_inches="tight")
    # plt.show()


if __name__ == "__main__":
    dataset_dir = "训练集合并_困惑度"

    datas = {
        "BERT+GAT+Seq2Seq": None,
        "BERT+Seq2Seq": None,
        "Seq2Seq": None,
        "GAT+Seq2Seq": None
    }

    dir_paths = []
    for dir_path, dir_names, filenames in os.walk(dataset_dir):
        dir_paths.append(dir_path)
    dir_paths = dir_paths[1:]

    for i, dir_path in enumerate(dir_paths):
        acc_file = os.path.join(dir_path, "test_acc.txt")
        config_file = os.path.join(dir_path, "config.txt")
        if i % 4 == 0:
            datas["BERT+GAT+Seq2Seq"] = moving_average_filter(_read_acc(acc_file), 12)
        elif i % 4 == 1:
            datas["BERT+Seq2Seq"] = moving_average_filter(_read_acc(acc_file), 12)
        elif i % 4 == 2:
            datas["Seq2Seq"] = moving_average_filter(_read_acc(acc_file), 12)
        elif i % 4 == 3:
            datas["GAT+Seq2Seq"] = moving_average_filter(_read_acc(acc_file), 12)

            with open(config_file, 'r') as f:
                config = json.load(f)
            save_file_name = "%d.svg" % (i // 4 + 1)
            save_file_name = os.path.join(dataset_dir, save_file_name)
            title = "seed=%d" % config["seed"]
            # title = ""
            draw_acc(datas, save_file=save_file_name, title=title)
            print("save file:", save_file_name)
